import cv2
import numpy as np
from typing import List, Optional, Any

from cv2 import Mat
from numpy import ndarray, dtype


class ImageStitcher:
    """图片拼接器类 - 直接执行拼接"""

    def __init__(self, mode: str = 'panorama'):
        """
        初始化拼接器

        Args:
            mode: 拼接模式 ('panorama', 'scans')
        """
        self.mode = mode
        if mode == 'panorama':
            self.stitcher = cv2.Stitcher.create(cv2.Stitcher_PANORAMA)
        elif mode == 'scans':
            self.stitcher = cv2.Stitcher.create(cv2.Stitcher_SCANS)
        else:
            raise ValueError("不支持的拼接模式，请选择 'panorama' 或 'scans'")

    def stitch_images(self, imgs: List[np.ndarray]) -> None | tuple[bool, ndarray | Any] | tuple[bool, None]:
        """
        拼接图片列表 - 直接返回拼接结果

        Args:
            imgs: 图片列表，要求已经是 CV_8UC3 格式

        Returns:
            拼接结果图片，如果失败返回None
        """
        if len(imgs) < 2:
            print("错误: 至少需要2张图片才能进行拼接")
            return None

        # print(f"开始拼接 {len(imgs)} 张图片...")

        # 直接进行拼接
        status, result = self.stitcher.stitch(imgs)

        if status == cv2.Stitcher_OK:
            print("拼接成功!")
            print(f"结果图片尺寸: {result.shape}")
            return True, result
        else:
            error_msg = self._get_error_message(status)
            # print(f'拼接失败: {error_msg}')
            return False, None


    def _get_error_message(self, status: int) -> str:
        """获取错误信息"""
        error_messages = {
            cv2.Stitcher_ERR_NEED_MORE_IMGS: "需要更多图片",
            cv2.Stitcher_ERR_HOMOGRAPHY_EST_FAIL: "单应性矩阵估计失败",
            cv2.Stitcher_ERR_CAMERA_PARAMS_ADJUST_FAIL: "相机参数调整失败"
        }
        return error_messages.get(status, f"未知错误 (错误代码: {status})")

    def save_result(self, result: np.ndarray, output_path: str) -> bool:
        """
        保存拼接结果

        Args:
            result: 拼接结果图片
            output_path: 保存路径

        Returns:
            是否保存成功
        """
        try:
            cv2.imwrite(output_path, result)
            print(f"结果已保存至: {output_path}")
            return True
        except Exception as e:
            print(f"保存失败: {str(e)}")
            return False

if __name__ == '__main__':
    imgs_path = ['temp_uploads\\image_4_1763004200.png', 'temp_uploads\\image_2_1763004200.png', 'temp_uploads\\image_3_1763004200.png', 'temp_uploads\\image_1_1763004200.png', 'temp_uploads\\image_0_1763004200.png']
    imgs = []
    for path in imgs_path:
        img = cv2.imread(path)
        imgs.append(img)
    stitch = ImageStitcher('panorama')
    success, result = stitch.stitch_images(imgs)
    result = cv2.resize(result, (1024,768))
    cv2.imshow('result', result)
    cv2.waitKey()
    cv2.destroyAllWindows()

